package com.langll.ai;

import com.alibaba.fastjson.JSONObject;
import com.langll.ai.model.ActorsFilms;
import com.langll.ai.tool.DateTimeTools;

import lombok.extern.slf4j.Slf4j;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.test.context.junit4.SpringRunner;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 *  测试聊天客户端
 */
@SpringBootTest(classes = AIChatApplication.class)
@RunWith(SpringRunner.class)
@Slf4j
public class TestChatClient {

    @Autowired
    private ChatClient chatClient;


    @Test
    public void generateCall(){
        // 创建一个 Prompt 对象，包含用户消息
        List<Message> messages = new ArrayList<>();
        String systemText = """
              将数据进行分类说明，总结数据变化情况
            """;

        SystemMessage systemMessage = new SystemMessage(systemText);
//        messages.add(systemMessage);
//        UserMessage userMessage = new UserMessage("现在是几号几点");
        UserMessage userMessage = new UserMessage("设置10分钟后的闹铃 ISO-8601 format");
        messages.add(userMessage);

        Prompt prompt = new Prompt(messages);
        String sessionId = UUID.randomUUID().toString();
        // 调用 chatClient 的 prompt 方法发送提示并获取响应
        ChatClient.CallResponseSpec response = chatClient.prompt(prompt)
//                .advisors(advisorSpec -> {
//                    advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, sessionId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100);
//                })
                .tools(new DateTimeTools())
                .call();

        System.out.println(response.content());
    }

    @Autowired
    ToolCallbackProvider tools;

    @Test
    public void functionCallback(){
        JSONObject addressJsonObject = new JSONObject();
        addressJsonObject.put("businessSegment", "Consumer");
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("request", addressJsonObject);

        log.info("-------------------{}", tools.getToolCallbacks().length);
        System.out.println("-------------------" + tools.getToolCallbacks().length);
        FunctionCallback[] functionCallbacks = tools.getToolCallbacks();
        for (FunctionCallback callback : functionCallbacks) {
            System.out.println(callback.getName());
            System.out.println("-------------------");
            System.out.println(callback.call(jsonObject.toJSONString()));
        }
    }

    @Test
    public void generateStream(){
        var converter = new BeanOutputConverter<>(new ParameterizedTypeReference<List<ActorsFilms>>() {});
        String t = converter.getFormat();
        Flux<String> flux = this.chatClient.prompt()
                .user(u -> u.text("""
                        Generate the filmography for a random actor.
                        {format}
                      """)
                        .param("format", converter.getFormat()))
                .stream()
                .content();
        String content = flux.collectList().block().stream().collect(Collectors.joining());
        List<ActorsFilms> actorFilms = converter.convert(content);
        System.out.println(actorFilms);
    }
}
